#! /usr/bin/env python3
# -*- coding: utf-8 -*-

import os
import argparse
import logging
import rosbag2_py
from rclpy.serialization import deserialize_message
from rosidl_runtime_py.utilities import get_message
import shutil
from datetime import datetime
import pytz
from typing import Tuple, Optional, List, Dict
from msg_register import Register  # Register类在msg_register.py中

# 配置日志格式
logging.basicConfig(
    level=logging.DEBUG,
    format="%(asctime)s - %(levelname)s - %(message)s"
)

class BagFilter(object):
    def __init__(self, 
                 input_bag: str, 
                 output_root: str, 
                 topics: Optional[list] = None,
                 msg_path: str = "",
                 use_source_mode: bool = False):
        self._input_bag = input_bag
        self._output_root = output_root  # 输出根目录
        self._topics = topics            # 需要保留的topic列表（可选）
        self._reader = rosbag2_py.SequentialReader()
        self._writer = rosbag2_py.SequentialWriter()

        # 注册自定义消息类型
        self._register_msg_types(msg_path, use_source_mode)
        
        # 截包条件配置
        self._start_topic = "/FM/FctReq"                  # 起始条件topic
        self._start_field = "/fct_data/apa_parking_state" # 起始条件字段
        self._start_value = 4                             # 起始值
        self._end_topic1 = "/PK/CUS/TargetSlot"           # 第一段终止topic
        self._end_field1 = "/slot_accuracy_visual"         # 第一段终止字段
        self._end_value1 = 1                              # 第一段终止值
        self._end_topic2 = "/FM/FctReq"                   # 第二段终止topic
        self._end_field2 = "/fct_data/apa_parking_state"  # 第二段终止字段
        self._end_value2 = 8                              # 第二段终止值

        # 存储截包片段（列表中的每个元素为一个片段的起止时间）
        self._segments: List[Dict[str, int]] = []
        # 存储bag的最后时间戳（用于处理未正常终止的情况）
        self._bag_end_ts: Optional[int] = None

    # 消息注册方法
    def _register_msg_types(self, msg_path: str, use_source_mode: bool):
        """注册自定义消息类型，确保解析时能识别字段"""
        if not msg_path:
            logging.warning("未指定自定义消息路径（--msg-path），可能导致字段解析失败")
            return
        
        # 验证消息路径存在
        if not os.path.isdir(msg_path):
            logging.error(f"自定义消息路径不存在：{msg_path}")
            return
        
        # 调用Register类注册消息
        try:
            register = Register()
            register.register_msg_types(msg_path, use_source_mode=use_source_mode)
            logging.info(f"自定义消息类型注册成功，路径：{msg_path}")
        except Exception as e:
            logging.error(f"消息注册失败：{str(e)}，可能导致截包异常")

    def _get_field_value(self, msg, field_path: str):
        # 统一把 / 换成 .
        field_path = field_path.replace('/', '.')
        # 去掉可能的第一个点
        field_path = field_path.lstrip('.')

        current = msg
        for field in field_path.split('.'):
            current = getattr(current, field, None)
            if current is None:
                return None
        return current

    def _get_message_timestamp(self, msg, bag_ts: int) -> int:
        """获取消息的时间戳（优先header.stamp，否则用bag时间戳）"""
        try:
            stamp = msg.header.stamp
            return stamp.sec * 10**9 + stamp.nanosec
        except AttributeError:
            logging.warning(f"消息无header.stamp，使用bag时间戳")
            return bag_ts

    def _ns_to_beijing_time(self, ns: int) -> str:
        """将纳秒时间戳转换为北京时间字符串（格式：HHMMSS，仅保留时分秒）"""
        # 转换纳秒到秒
        seconds = ns // 10**9
        # 创建UTC时间对象
        utc_dt = datetime.utcfromtimestamp(seconds)
        # 转换为北京时间（UTC+8）
        beijing_tz = pytz.timezone('Asia/Shanghai')
        beijing_dt = utc_dt.replace(tzinfo=pytz.utc).astimezone(beijing_tz)
        # 格式化输出（仅保留时分秒，不含年月日）
        return beijing_dt.strftime("%H%M%S")

    def scan_segments(self):
        """扫描bag，获取所有满足条件的截包片段"""
        reader = rosbag2_py.SequentialReader()
        in_storage_options = rosbag2_py._storage.StorageOptions(
            uri=self._input_bag, 
            storage_id='sqlite3'
        )
        reader.open(in_storage_options, rosbag2_py._storage.ConverterOptions('', ''))

        # 获取所有topic类型映射
        topic_types = reader.get_all_topics_and_types()
        type_map = {t.name: t.type for t in topic_types}

        # 验证必要的topic是否存在
        required_topics = [self._start_topic]  # 只强制检查起始topic，其他topic不存在时降级处理
        for topic in required_topics:
            if topic not in type_map:
                logging.error(f"必要topic {topic} 不在bag中，无法进行截包")
                return False

        # 过滤掉无关 topic
        needed = {
            self._start_topic,
            self._end_topic1,
            self._end_topic2,
            "/SA/INSPVA",
            "/PK/DR/Result",
            "/VA/VehicleStatusIpd",
            "/VA/VehicleMotionIpd"
        }
        storage_filter = rosbag2_py.StorageFilter(
            topics=[t for t in needed if t in type_map]
        )
        reader.set_filter(storage_filter)

        # 状态变量
        current_segment = None  # 当前片段状态：{start_ts, stage}，stage=1表示等待第一段结束，stage=2表示等待第二段结束
        prev_start_value = None
        prev_end1_value = None
        prev_end2_value = None
        self._bag_end_ts = 0  # 初始化bag结束时间戳

        logging.info("开始扫描bag，寻找截包片段...")
        while reader.has_next():
            topic, data, bag_ts = reader.read_next()
            # 更新bag结束时间戳
            if bag_ts > self._bag_end_ts:
                self._bag_end_ts = bag_ts
                
            msg_type = get_message(type_map[topic])
            msg = deserialize_message(data, msg_type)
            msg_ts = self._get_message_timestamp(msg, bag_ts)

            # 处理起始条件：如果start_field是4，且尚未开始截取，则立即开始
            # logging.info(f"topic: {topic}, msg_ts: {msg_ts}, bag_ts: {bag_ts}")
            if topic == self._start_topic:
                current_value = self._get_field_value(msg, self._start_field)
                # logging.info(f"topic: {topic}, current_value: {current_value}, beijing_time: {self._ns_to_beijing_time(msg_ts)}")
                if current_value is None:
                    continue
                    
                # 起始条件修改：只要字段值为4且当前没有活跃片段，就开始截取
                if current_value == self._start_value and current_segment is None:
                    current_segment = {
                        "start_ts": msg_ts,
                        "stage": 1  # 处于第一段（重规划前）
                    }
                    # 转换为北京时间并记录
                    beijing_time = self._ns_to_beijing_time(msg_ts)
                    logging.info(
                        f"检测到片段起始点：{self._start_field}为{self._start_value}，"
                        f"北京时间（时分秒）：{beijing_time}，时间戳：{msg_ts} ns"
                    )
                prev_start_value = current_value

            # 处理第一段终止条件：end1_field从非1变为1（仅在stage=1时生效）
            if (topic == self._end_topic1 and 
                current_segment is not None and 
                current_segment["stage"] == 1):
                current_value = self._get_field_value(msg, self._end_field1)
                if current_value is None:
                    continue
                    
                if (prev_end1_value != self._end_value1 and 
                    current_value == self._end_value1):
                    # 检测到第一段终止条件：从其他值变为1
                    current_segment["end1_ts"] = msg_ts
                    current_segment["stage"] = 2  # 进入第二段（重规划后）
                    # 转换为北京时间并记录
                    beijing_time = self._ns_to_beijing_time(msg_ts)
                    logging.info(
                        f"检测到第一段终止点：{self._end_field1}从{prev_end1_value}变为{self._end_value1}，"
                        f"北京时间（时分秒）：{beijing_time}，时间戳：{msg_ts} ns"
                    )
                prev_end1_value = current_value

            # 处理第二段终止条件：end2_field从非8变为8（仅在stage=2时生效）
            if (topic == self._end_topic2 and 
                current_segment is not None):
                current_value = self._get_field_value(msg, self._end_field2)
                if current_value is None:
                    continue
                    
                # 对于处于stage=1的片段，也监控end2_value，用于处理异常情况
                if (prev_end2_value != self._end_value2 and 
                    current_value == self._end_value2):
                    # 检测到第二段终止条件：从其他值变为8
                    current_segment["end2_ts"] = msg_ts
                    # 转换为北京时间并记录
                    beijing_time = self._ns_to_beijing_time(msg_ts)
                    
                    # 根据当前阶段判断如何处理
                    if current_segment["stage"] == 1:
                        # 未触发第一段终止就直接到第二段终止，只生成一个包
                        logging.warning(
                            f"未检测到第一段终止条件，直接到第二段终止，仅生成一个完整包，"
                            f"北京时间（时分秒）：{beijing_time}"
                        )
                    else:
                        # 正常流程，生成两个包
                        logging.info(
                            f"检测到第二段终止点：{self._end_field2}从{prev_end2_value}变为{self._end_value2}，"
                            f"北京时间（时分秒）：{beijing_time}，时间戳：{msg_ts} ns"
                        )
                    
                    self._segments.append(current_segment)
                    current_segment = None  # 重置，准备下一个片段
                prev_end2_value = current_value

        # 处理扫描结束时仍未终止的片段
        if current_segment is not None:
            # 使用bag的最后时间戳作为终止时间
            current_segment["end2_ts"] = self._bag_end_ts
            beijing_time = self._ns_to_beijing_time(self._bag_end_ts)
            logging.warning(
                f"未检测到第二段终止条件，使用bag结束时间作为终止点：{beijing_time}（北京时间，时分秒）"
            )
            
            # 检查是否有第一段终止
            if "end1_ts" not in current_segment:
                # 没有第一段终止，只生成一个包
                logging.warning(f"未检测到第一段终止条件，仅生成一个完整包")
            
            self._segments.append(current_segment)

        if not self._segments:
            logging.error("未找到任何满足条件的截包片段")
            return False

        logging.info(f"扫描完成，共找到 {len(self._segments)} 个片段")
        return True

    def filter_segments(self):
        """根据扫描到的片段截取并保存多个bag"""
        # 第一步：扫描片段
        if not self.scan_segments():
            return
        logging.info(f"片段列表：{self._segments}")

        # 第二步：打开输入bag
        in_storage_options = rosbag2_py._storage.StorageOptions(
            uri=self._input_bag, 
            storage_id='sqlite3'
        )
        self._reader.open(in_storage_options, rosbag2_py._storage.ConverterOptions('', ''))
        
        needed = {
            "/FM/FctReq",
            "/PK/CUS/TargetSlot",
            "/SA/INSPVA",
            "/PK/DR/Result",
            "/VA/VehicleStatusIpd",
            "/VA/VehicleMotionIpd"
        }
        storage_filter = rosbag2_py.StorageFilter(
            topics=[t for t in needed if t in {m.name for m in self._reader.get_all_topics_and_types()}])
        self._reader.set_filter(storage_filter)

        logging.info(f"过滤的 topic 列表：{needed}")
        # 第三步：为每个片段创建输出bag
        for i, segment in enumerate(self._segments, 1):
            # 转换时间戳为北京时间字符串（仅时分秒）
            start_time_str = self._ns_to_beijing_time(segment['start_ts'])
            
            # 检查是否有第一段终止
            if "end1_ts" in segment and segment["stage"] >= 2:
                # 有第一段终止，生成两个包
                end1_time_str = self._ns_to_beijing_time(segment['end1_ts'])
                end2_time_str = self._ns_to_beijing_time(segment['end2_ts'])
                
                # 重规划前：起始时间-第一段终止时间_before
                out_dir_before = os.path.join(
                    self._output_root, 
                    f"{start_time_str}-{end1_time_str}_before_replan"
                )
                # 重规划后：第一段终止时间-第二段终止时间_after
                out_dir_after = os.path.join(
                    self._output_root, 
                    f"{end1_time_str}-{end2_time_str}_after_replan"
                )
                
                # 截取重规划前的片段
                logging.info(f"\n开始截取重规划前片段：{start_time_str} - {end1_time_str}")
                self._write_segment(out_dir_before, segment['start_ts'], segment['end1_ts'])
                
                # 截取重规划后的片段
                logging.info(f"开始截取重规划后片段：{end1_time_str} - {end2_time_str}")
                self._write_segment(out_dir_after, segment['end1_ts'], segment['end2_ts'])
            else:
                # 没有第一段终止，只生成一个完整包
                end_time_str = self._ns_to_beijing_time(segment['end2_ts'])
                # 完整包：起始时间-终止时间_full
                out_dir_full = os.path.join(
                    self._output_root, 
                    f"{start_time_str}-{end_time_str}_no_replan"
                )
                logging.info(f"\n开始截取完整片段：{start_time_str} - {end_time_str}")
                self._write_segment(out_dir_full, segment['start_ts'], segment['end2_ts'])

        logging.info(f"\n所有 {len(self._segments)} 个片段截取完成！")
        logging.info(f"输出路径：{self._output_root}")

    def _write_segment(self, output_dir: str, start_ts: int, end_ts: int):
        """将指定时间范围内的消息写入输出bag"""

        # 1. 清理旧目录
        if os.path.exists(output_dir):
            shutil.rmtree(output_dir)

        # 2. 重新打开 reader（确保从头开始）
        reader = rosbag2_py.SequentialReader()
        in_storage_options = rosbag2_py._storage.StorageOptions(
            uri=self._input_bag, 
            storage_id='sqlite3'
        )
        reader.open(in_storage_options, rosbag2_py._storage.ConverterOptions('', ''))

        # 3. 获取所有 topic 类型映射
        topic_types = reader.get_all_topics_and_types()
        type_map = {t.name: t.type for t in topic_types}

        # 4. 设置 topic 过滤条件
        needed = {
            self._start_topic,
            self._end_topic1,
            self._end_topic2,
            "/SA/INSPVA",
            "/PK/DR/Result",
            "/VA/VehicleStatusIpd",
            "/VA/VehicleMotionIpd"
        }
        storage_filter = rosbag2_py.StorageFilter(
            topics=[t for t in needed if t in type_map]
        )
        reader.set_filter(storage_filter)

        # 5. 创建 writer（一步到位）
        writer = rosbag2_py.SequentialWriter()
        storage_options = rosbag2_py._storage.StorageOptions(
            uri=output_dir, storage_id='sqlite3')
        converter_options = rosbag2_py._storage.ConverterOptions('', '')
        writer.open(storage_options, converter_options)

        # 6. 注册所有 topic
        for topic_metadata in reader.get_all_topics_and_types():
            topic_info = rosbag2_py._storage.TopicMetadata(
                name=topic_metadata.name,
                type=topic_metadata.type,
                serialization_format='cdr')
            writer.create_topic(topic_info)

        # 7. 写消息
        msg_count = 0
        while reader.has_next():
            topic_name, data, t = reader.read_next()
            if start_ts <= t <= end_ts:
                writer.write(topic_name, data, t)
                msg_count += 1

        # 8. 释放
        del writer

        # 9. 可选验证
        metadata_path = os.path.join(output_dir, "metadata.yaml")
        if os.path.isfile(metadata_path):
            out_metadata = rosbag2_py.Info().read_metadata(output_dir, 'sqlite3')
            logging.info(
                f"片段 {output_dir} 写入完成，消息数：{msg_count}，"
                f"时长：{out_metadata.duration.nanoseconds / 1e9:.2f} 秒"
            )
        else:
            logging.warning("metadata.yaml 未生成，跳过时长读取")


def delete_dir_files(output_dir: str):
    """删除输出目录（若存在）"""
    if os.path.exists(output_dir):
        shutil.rmtree(output_dir)
        logging.info(f"已删除现有目录：{output_dir}")


def main():
    # 获取脚本所在目录，用于设置默认路径（改为上级目录）
    script_dir = os.path.dirname(os.path.abspath(__file__))
    parent_dir = os.path.dirname(script_dir)  # 获取上级目录
    default_input = os.path.join(parent_dir, "bag")  # 默认输入路径：上级目录bag文件夹
    default_output = os.path.join(parent_dir, "truncated")  # 默认输出路径：上级目录truncated文件夹
    default_msg_path = os.path.join(parent_dir, "install")  # 默认ros安装路径：上级目录install文件夹

    parser = argparse.ArgumentParser(description="根据topic字段值自动截取ROS2 bag（重规划前后片段）")
    # 路径参数（支持默认值）
    parser.add_argument(
        "-i", "--input_bag", 
        type=str, 
        default=default_input, 
        help=f"输入bag文件或文件夹路径（默认：{default_input}）"
    )
    parser.add_argument(
        "-o", "--output_dir", 
        type=str, 
        default=default_output, 
        help=f"输出根目录（默认：{default_output}）"
    )
    parser.add_argument(
        "-m", "--msg_path", 
        type=str, 
        default=default_msg_path,
        help=f"自定义消息(.msg文件)的根目录路径（默认：{default_msg_path}）"
    )
    # 可选参数
    parser.add_argument(
        "-k", "--keep_topics", 
        nargs="+", 
        type=str, 
        help="需要保留的topic列表（不指定则保留所有）"
    )
    # 在命令行参数中添加选项
    parser.add_argument(
        "--use-source-mode",
        action="store_true",
        help="使用ROS2 msg源文件模式（默认使用install目录模式）"
    )
    args = parser.parse_args()

    # 处理输入路径：如果是文件夹，自动查找其中的.db3文件
    input_path = args.input_bag
    if os.path.isdir(input_path):
        # 查找文件夹中的.db3文件
        db3_files = [f for f in os.listdir(input_path) if f.endswith(".db3")]
        if not db3_files:
            logging.error(f"输入文件夹 {input_path} 中未找到.db3文件")
            return
        if len(db3_files) > 1:
            logging.warning(f"输入文件夹 {input_path} 中有多个.db3文件，使用第一个：{db3_files[0]}")
        input_bag = os.path.join(input_path, db3_files[0])
    else:
        input_bag = input_path

    # 验证输入文件存在
    if not os.path.isfile(input_bag):
        logging.error(f"输入bag文件不存在：{input_bag}")
        return

    # 准备输出目录
    output_root = args.output_dir
    os.makedirs(output_root, exist_ok=True)
    logging.info(f"输出根目录：{output_root}")

    # 执行截包
    filter = BagFilter(
        input_bag=input_bag,
        output_root=output_root,
        topics=args.keep_topics,
        msg_path=args.msg_path
    )
    filter.filter_segments()


if __name__ == "__main__":
    main()
    